#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/Index/IR/IndexOps.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/X86Vector/X86VectorDialect.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypeInterfaces.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/IR/TypeRange.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Pass/PassRegistry.h"
#include "mlir/Support/LogicalResult.h"
#include "mlir/Support/TypeID.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include <iostream>
#include <ostream>

#include "crypto/Crypto.h"
#include "llvm/ADT/ArrayRef.h"

namespace mlir {
#define GEN_PASS_DEF_CONVERTCRYPTOTOARM
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir

using namespace mlir;
using namespace vector;

//===----------------------------------------------------------------------===//
// Operation conversion
//===----------------------------------------------------------------------===//

namespace {
class Sha1LoweringFPGAPattern : public OpRewritePattern<crypto::Sha1_OP> {
public:
  using OpRewritePattern<crypto::Sha1_OP>::OpRewritePattern;
  LogicalResult matchAndRewrite(crypto::Sha1_OP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();

    // create extern function sha1_run if not exist.
    StringRef algo_func_name = "sha1_run";
    auto m = op->getParentOfType<ModuleOp>();
    func::FuncOp existingFunc = m.lookupSymbol<func::FuncOp>(algo_func_name);
    if (!existingFunc) {
      auto algo_run_ty = rewriter.getFunctionType(
          mlir::TypeRange{
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              rewriter.getIntegerType(32),
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8))},
          mlir::TypeRange{});
      auto algo_run = func::FuncOp::create(loc, "sha1_run", algo_run_ty);
      algo_run.setVisibility(mlir::SymbolTable::Visibility::Private);
      m.push_back(algo_run);
    }

    auto funcOp = rewriter.create<func::CallOp>(
        loc, algo_func_name, mlir::TypeRange{},
        ValueRange{op.getA(), op.getB(), op.getC()});
    rewriter.replaceOp(op, funcOp->getResults());
    // rewriter.eraseOp(op);
    return mlir::success();
  }
};

class Sha1CrackLoweringFPGAPattern
    : public OpRewritePattern<crypto::SHA1CrackOp> {
public:
  using OpRewritePattern<crypto::SHA1CrackOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(crypto::SHA1CrackOp op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();

    // create extern function sha1Crack_run if not exist.
    StringRef algo_func_name = "sha1Crack_run";
    auto m = op->getParentOfType<ModuleOp>();
    func::FuncOp existingFunc = m.lookupSymbol<func::FuncOp>(algo_func_name);
    if (!existingFunc) {
      auto algo_run_ty = rewriter.getFunctionType(
          mlir::TypeRange{
              // hash
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              // input
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              // MSG_NUM
              rewriter.getIntegerType(32),
              // MSG_LEN
              rewriter.getIntegerType(32)},
          mlir::TypeRange{
              // OUTPUT
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8))});
      auto algo_run = func::FuncOp::create(loc, "sha1Crack_run", algo_run_ty);
      algo_run.setVisibility(mlir::SymbolTable::Visibility::Private);
      m.push_back(algo_run);
    }

    auto funcOp = rewriter.create<func::CallOp>(
        loc, algo_func_name, mlir::TypeRange{op.getRes().getType()},
        ValueRange{op.getTarget(), op.getCandidates(), op.getNumMessages(),
                   op.getMessageLen()});
    rewriter.replaceOp(op, funcOp->getResults());
    // rewriter.eraseOp(op);
    return mlir::success();
  }
};

// =============
// MD5
// =============
class Md5LoweringFPGAPattern : public OpRewritePattern<crypto::MD5OP> {
public:
  using OpRewritePattern<crypto::MD5OP>::OpRewritePattern;
  LogicalResult matchAndRewrite(crypto::MD5OP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();

    // create extern function md5_run if not exist.
    StringRef algo_func_name = "md5_run";
    auto m = op->getParentOfType<ModuleOp>();
    func::FuncOp existingFunc = m.lookupSymbol<func::FuncOp>(algo_func_name);
    if (!existingFunc) {
      auto algo_run_ty = rewriter.getFunctionType(
          mlir::TypeRange{
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              rewriter.getIntegerType(32),
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8))},
          mlir::TypeRange{});
      auto algo_run = func::FuncOp::create(loc, "md5_run", algo_run_ty);
      algo_run.setVisibility(mlir::SymbolTable::Visibility::Private);
      m.push_back(algo_run);
    }

    auto funcOp = rewriter.create<func::CallOp>(
        loc, algo_func_name, mlir::TypeRange{},
        ValueRange{op.getInput(), op.getMsgLen(), op.getOutput()});
    rewriter.replaceOp(op, funcOp->getResults());
    // rewriter.eraseOp(op);
    return mlir::success();
  }
};

class Md5CrackLoweringFPGAPattern
    : public OpRewritePattern<crypto::MD5CrackOp> {
public:
  using OpRewritePattern<crypto::MD5CrackOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(crypto::MD5CrackOp op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();

    // create extern function md5_run if not exist.
    StringRef algo_func_name = "md5Crack_run";
    auto m = op->getParentOfType<ModuleOp>();
    func::FuncOp existingFunc = m.lookupSymbol<func::FuncOp>(algo_func_name);
    if (!existingFunc) {
      auto algo_run_ty = rewriter.getFunctionType(
          mlir::TypeRange{
              // hash
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              // input
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              // MSG_NUM
              rewriter.getIntegerType(32),
              // MSG_LEN
              rewriter.getIntegerType(32)},
          mlir::TypeRange{
              // OUTPUT
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8))});
      auto algo_run = func::FuncOp::create(loc, "md5Crack_run", algo_run_ty);
      algo_run.setVisibility(mlir::SymbolTable::Visibility::Private);
      m.push_back(algo_run);
    }

    auto funcOp = rewriter.create<func::CallOp>(
        loc, algo_func_name, mlir::TypeRange{op.getRes().getType()},
        ValueRange{op.getTarget(), op.getCandidates(), op.getNumMessages(),
                   op.getMessageLen()});
    rewriter.replaceOp(op, funcOp->getResults());
    // rewriter.eraseOp(op);
    return mlir::success();
  }
};

// ================
// SM3
// ================
class SM3LoweringFPGAPattern : public OpRewritePattern<crypto::SM3OP> {
public:
  using OpRewritePattern<crypto::SM3OP>::OpRewritePattern;
  LogicalResult matchAndRewrite(crypto::SM3OP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();

    // create extern function sm3_run if not exist.
    StringRef algo_func_name = "sm3_run";
    auto m = op->getParentOfType<ModuleOp>();
    func::FuncOp existingFunc = m.lookupSymbol<func::FuncOp>(algo_func_name);
    if (!existingFunc) {
      auto algo_run_ty = rewriter.getFunctionType(
          mlir::TypeRange{
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              rewriter.getIntegerType(32),
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8))},
          mlir::TypeRange{});
      auto algo_run = func::FuncOp::create(loc, "sm3_run", algo_run_ty);
      algo_run.setVisibility(mlir::SymbolTable::Visibility::Private);
      m.push_back(algo_run);
    }

    auto funcOp = rewriter.create<func::CallOp>(
        loc, algo_func_name, mlir::TypeRange{},
        ValueRange{op.getInput(), op.getMsgLen(), op.getOutput()});
    rewriter.replaceOp(op, funcOp->getResults());
    // rewriter.eraseOp(op);
    return mlir::success();
  }
};

class SM3CrackLoweringFPGAPattern
    : public OpRewritePattern<crypto::SM3CrackOp> {
public:
  using OpRewritePattern<crypto::SM3CrackOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(crypto::SM3CrackOp op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();

    // create extern function sm3Crack_run if not exist.
    StringRef algo_func_name = "sm3Crack_run";
    auto m = op->getParentOfType<ModuleOp>();
    func::FuncOp existingFunc = m.lookupSymbol<func::FuncOp>(algo_func_name);
    if (!existingFunc) {
      auto algo_run_ty = rewriter.getFunctionType(
          mlir::TypeRange{
              // hash
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              // input
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              // MSG_NUM
              rewriter.getIntegerType(32),
              // MSG_LEN
              rewriter.getIntegerType(32)},
          mlir::TypeRange{
              // OUTPUT
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8))});
      auto algo_run = func::FuncOp::create(loc, "sm3Crack_run", algo_run_ty);
      algo_run.setVisibility(mlir::SymbolTable::Visibility::Private);
      m.push_back(algo_run);
    }

    auto funcOp = rewriter.create<func::CallOp>(
        loc, algo_func_name, mlir::TypeRange{op.getRes().getType()},
        ValueRange{op.getTarget(), op.getCandidates(), op.getNumMessages(),
                   op.getMessageLen()});
    rewriter.replaceOp(op, funcOp->getResults());
    // rewriter.eraseOp(op);
    return mlir::success();
  }
};

// ================
// AES
// ================
class AES128LoweringFPGAPattern : public OpRewritePattern<crypto::AESOP> {
public:
  using OpRewritePattern<crypto::AESOP>::OpRewritePattern;
  LogicalResult matchAndRewrite(crypto::AESOP op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();

    // create extern function sm3_run if not exist.
    StringRef algo_func_name = "aes_run";
    auto m = op->getParentOfType<ModuleOp>();
    func::FuncOp existingFunc = m.lookupSymbol<func::FuncOp>(algo_func_name);
    if (!existingFunc) {
      auto algo_run_ty = rewriter.getFunctionType(
          mlir::TypeRange{
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              rewriter.getIntegerType(32),
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
          },
          mlir::TypeRange{});
      auto algo_run = func::FuncOp::create(loc, "aes_run", algo_run_ty);
      algo_run.setVisibility(mlir::SymbolTable::Visibility::Private);
      m.push_back(algo_run);
    }

    auto funcOp = rewriter.create<func::CallOp>(
        loc, algo_func_name, mlir::TypeRange{},
        ValueRange{op.getInput(), op.getMsgLen(), op.getOutput(), op.getKey(),
                   op.getSBoxBase()});
    rewriter.replaceOp(op, funcOp->getResults());
    // rewriter.eraseOp(op);
    return mlir::success();
  }
};

class AES128CrackLoweringFPGAPattern
    : public OpRewritePattern<crypto::AESCrackOp> {
public:
  using OpRewritePattern<crypto::AESCrackOp>::OpRewritePattern;
  LogicalResult matchAndRewrite(crypto::AESCrackOp op,
                                PatternRewriter &rewriter) const override {
    auto loc = op.getLoc();

    // create extern function aesCrack_run if not exist.
    StringRef algo_func_name = "aesCrack_run";
    auto m = op->getParentOfType<ModuleOp>();
    func::FuncOp existingFunc = m.lookupSymbol<func::FuncOp>(algo_func_name);
    if (!existingFunc) {
      auto algo_run_ty = rewriter.getFunctionType(
          mlir::TypeRange{
              // input
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              // msgLen
              rewriter.getIntegerType(32),
              // outputs
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              // cptLen
              rewriter.getIntegerType(32),
              // key
              mlir::MemRefType::get(
                  llvm::ArrayRef<int64_t>{mlir::ShapedType::kDynamic},
                  rewriter.getIntegerType(8)),
              // keyLen
              rewriter.getIntegerType(32),
              // msgCnt
              rewriter.getIntegerType(32),
          },
          mlir::TypeRange{});
      auto algo_run = func::FuncOp::create(loc, "aesCrack_run", algo_run_ty);
      algo_run.setVisibility(mlir::SymbolTable::Visibility::Private);
      m.push_back(algo_run);
    }

    auto funcOp = rewriter.create<func::CallOp>(
        loc, algo_func_name, mlir::TypeRange{},
        ValueRange{op.getInputs(), op.getMsgLen(), op.getOutputs(),
                   op.getCptLen(), op.getKey(), op.getKeyLen(),
                   op.getMsgCnt()});
    rewriter.replaceOp(op, funcOp->getResults());
    // rewriter.eraseOp(op);
    return mlir::success();
  }
};

} // namespace

void populateConvertCryptoToFPGAPatterns(RewritePatternSet &patterns) {
  patterns.add<Sha1LoweringFPGAPattern, Md5LoweringFPGAPattern,
               SM3LoweringFPGAPattern, AES128LoweringFPGAPattern,
               Sha1CrackLoweringFPGAPattern, Md5CrackLoweringFPGAPattern,
               SM3CrackLoweringFPGAPattern, AES128CrackLoweringFPGAPattern>(
      patterns.getContext());
}

namespace {
class LowerCryptoToFPGAPass
    : public PassWrapper<LowerCryptoToFPGAPass, OperationPass<ModuleOp>> {
public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerCryptoToFPGAPass)
  LowerCryptoToFPGAPass() = default;
  LowerCryptoToFPGAPass(const LowerCryptoToFPGAPass &) {}
  StringRef getArgument() const final { return "lower-crypto-to-fpga"; }
  StringRef getDescription() const final {
    return "Lower Crypto Dialect to FPGA function call.";
  }

  void runOnOperation() override;
  void getDependentDialects(DialectRegistry &registry) const override {
    registry.insert<func::FuncDialect>();
  }
};
} // end anonymous namespace.

void LowerCryptoToFPGAPass::runOnOperation() {
  MLIRContext *context = &getContext();
  ModuleOp module = getOperation();

  ConversionTarget target(*context);
  target.addLegalDialect<func::FuncDialect, memref::MemRefDialect,
                         arith::ArithDialect>();
  target.addLegalOp<ModuleOp, func::FuncOp, func::ReturnOp>();

  RewritePatternSet patterns(context);
  populateConvertCryptoToFPGAPatterns(patterns);

  if (failed(applyPartialConversion(module, target, std::move(patterns))))
    signalPassFailure();
}

void registerLowerCryptoToFPGAPass() {
  PassRegistration<LowerCryptoToFPGAPass>();
}

std::unique_ptr<Pass> createLowerCryptoToFPGAPass() {
  return std::make_unique<LowerCryptoToFPGAPass>();
}
